# 研究机构：上海大学机自学院
# 研 究 生：王强
# 开发时间：2024/6/3 10:48
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import math

font = {'family': 'SimSun',  # 宋体
        # 'weight': 'bold',  # 加粗
        'size': '10.5'  # 五号
        }
plt.rc('font', **font)
plt.rc('axes', unicode_minus=False)

class EKF:
    def __init__(self, px, vx, py, vy, w, T):
        # 本例中，状态x为（坐标x， 坐标y， 速度x， 速度y， 雷达采样周期T），观测值z为（坐标x， 坐标y）
        self.B = 0                                  # 外部输入为0
        self.u = 0                                  # 外部输入为0
        self.K = float('nan')                       # 卡尔曼增益无需初始化
        self.z = float('nan')                       # 这里无需初始化，每次使用kf_update之前需要输入观察值z
        self.P = np.diag(np.ones(5))*0.01           # 初始P设为0 ??? zeros(4, 4)
        self.S = np.diag(np.ones(2))
        # 初始状态：函数外部提供初始化的状态，本例使用观察值进行初始化，vx，vy初始为0
        self.x = [px, vx, py, vy, w]
        self.pre_nextz = [0, 0]
        self.target_speed = 100000                       # 单位字节 1节 = 0.5144 m/s
        self.taeget_course = 1000000                      # -180~180
        self.target_distance = 10000000                    # 海里
        self.target_bearing = 100000                            # 方位角 -180~180
        # 状态转移矩阵A
        # 和线性系统的预测机制有关，这里的线性系统是上一刻的位置加上速度等于当前时刻的位置，而速度本身保持不变
        self.A = np.array([[1, np.sin(w * T) / w, 0, -(1 - np.cos(w * T)) / w, 0],
                           [0, np.cos(w * T), 0, -np.sin(w * T), 0],
                           [0, (1 - np.cos(w * T)) / w, 1, np.sin(w * T) / w, 0],
                           [0, np.sin(w * T), 0, np.cos(w * T), 0],
                           [0, 0, 0, 0, 1]])
        q_v = np.diag([3 ** 2, 3 ** 2, 0.001]) * 0.5
        G_v = np.array([[T ** 2 / 2, 0, 0],
                        [T, 0, 0],
                        [0, T ** 2 / 2, 0],
                        [0, T, 0],
                        [0, 0, T]])
        self.Q = np.dot(np.dot(G_v, q_v), np.transpose(G_v))
        # 预测噪声协方差矩阵Q：假设预测过程上叠加一个高斯噪声，协方差矩阵为Q
        # 大小取决于对预测过程的信任程度。比如，假设认为运动目标在y轴上的速度可能不匀速，那么可以把这个对角矩阵
        # 的最后一个值调大。有时希望出来的轨迹更平滑，可以把这个调更小
        # 观测矩阵H：z = H * x
        # 这里的状态是（坐标x， 坐标y， 速度x， 速度y），观察值是（坐标x， 坐标y），所以H = eye(2, 4)
        self.H = np.array([[1, 0, 0, 0, 0],
                           [0, 0, 1, 0, 0]])

        # 观测噪声协方差矩阵R：假设观测过程上存在一个高斯噪声，协方差矩阵为R
        # 大小取决于对观察过程的信任程度。比如，假设观测结果中的坐标x值常常很准确，那么矩阵R的第一个值应该比较小
        self.R = np.diag([10**2, 10**2])

    def KF_update(self):
        # 以下为卡尔曼滤波的五个方程（步骤）
        a1 = np.dot(self.A, self.x)  # np.dot 表示为矩阵相乘，即A*B A的列数等于B的行数
        a2 = self.B * self.u
        x_ = np.array(a1) + np.array(a2)  # x_ 为卡尔曼滤波的预测值 即x_ = A*x +B*u 这里B=U=0

        b1 = np.dot(self.A, self.P)
        b2 = np.dot(b1, np.transpose(self.A))  # np.transpose 表示为矩阵的转置
        p_ = np.array(b2) + np.array(self.Q)  # p_ 为先验误差协方差方程矩阵 p_ = A*P*A^T + Q 这里P为上一时刻后验误差协方差矩阵

        c1 = np.dot(p_, np.transpose(self.H))
        c2 = np.dot(self.H, p_)
        c3 = np.dot(c2, np.transpose(self.H))

        self.S = np.array(c3) + np.array(self.R)
        c5 = np.linalg.matrix_power(self.S, -1)  # np.linalg.matrix_power 表示为矩阵的求逆
        self.K = np.dot(c1, c5)  # K 为卡尔曼增益 K = p_k*H^T*(H*p_*H^T+R)^-1 这里 H 为观测转移矩阵

        d1 = np.dot(self.H, x_)
        # d1 = [np.linalg.norm([x_[0], x_[2]], ord=2), math.atan2(x_[2], x_[0])/3.14*180]
        d2 = np.array(self.z) - np.array(d1)
        d3 = np.dot(self.K, d2)
        self.x = np.array(x_) + np.array(d3)  # x 为当前时刻的最优估计值 x = x_ + K(z-H*x_)

        e1 = np.dot(self.K, self.H)
        e2 = np.dot(e1, p_)
        self.P = np.array(p_) - np.array(e2)  # P 为后验误差协方差矩阵 P = (p_ - K*H*p_)

        # 预测下一时刻的值，用于全局最近邻数据关联算法
        f1 = np.dot(self.A, self.x)
        self.pre_nextz = [np.linalg.norm([f1[0], f1[2]], ord=2), math.atan2(f1[2], f1[0]) / 3.14 * 180]
        # 速度单位为节 1节 = 0.5144 m/s
        self.target_speed = np.linalg.norm([self.x[1], self.x[3]], ord=2) / 0.51444444
        # 航向
        self.taeget_course = math.atan2(self.x[3], self.x[1]) / 3.14 * 180
        # 距离 1海里=1.852公里(千米)
        self.target_distance = np.linalg.norm([self.x[0], self.x[2]], ord=2) / 1852
        # 方位角
        self.target_bearing = math.atan2(self.x[2], self.x[0]) / 3.14 * 180
if __name__ == '__main__':
    # 真实路径
    path = './X_ideal.xlsx'
    data_A = pd.read_excel(path, header=None, skiprows=[0])
    data_A_x = list(data_A.iloc[::, 0])
    data_A_y = list(data_A.iloc[::, 1])
    A = np.array(list(zip(data_A_x, data_A_y)))

    # plt.subplot(131)
    plt.figure()
    plt.plot(data_A_x, data_A_y, 'b-+')
    # plt.title('实际的真实路径')

    # 检测到的路径
    path = './Z_obs.xlsx'
    data_B = pd.read_excel(path, header=None, skiprows=[0])
    radius = list(data_B.iloc[::, 0])
    azimuth = list(data_B.iloc[::, 1])
    B = np.array(list(zip(radius, azimuth)))
    x_obs = radius * np.cos(np.array(azimuth)/180*3.14)
    y_obs = radius * np.sin(np.array(azimuth)/180*3.14)
    print(y_obs)
    # plt.subplot(132)
    plt.plot(x_obs, y_obs, 'r-+')
    # plt.title('检测到的路径')

    # 卡尔曼滤波
    kf_params_record = np.zeros((len(data_B), 5))
    kf_params_p = np.zeros((len(data_B), 5))
    t = len(data_B)
    kalman_filter = EKF(px=x_obs[0], vx=0, py=y_obs[0], vy=0, w=0.01, T=1)  # 初始化
    for i in range(t):
            # print([data_B_x[i], data_B_y[i]])
        kalman_filter.z = np.transpose([x_obs[i], y_obs[i]])  # 设置当前时刻的观测位置
        kalman_filter.KF_update()  # 卡尔曼滤波
        kf_params_record[i, ::] = np.transpose(kalman_filter.x)
        # kf_params_p[i, ::] = np.transpose(kalman_filter_params.G)

    kf_trace = kf_params_record[::, [0, 2]]
    # kf_trace_1 = kf_params_p[::, :2]

    # plt.subplot(133)
    plt.plot(kf_trace[::, 0], kf_trace[::, 1], 'g-+')
    # plt.plot(kf_trace_1[1:26, 0], kf_trace_1[1:26, 1], 'm-+')
    legend = ['实际的真实路径', '检测路径', '预测路径']
    plt.legend(legend, loc="best", frameon=False)
    plt.title('卡尔曼滤波后的效果')
    plt.savefig('result.svg', dpi=600)
    plt.show()
    # plt.close()

    # p = accuracy(kf_trace, A)
    # print(p)